import asyncio
from pyVim.connect import SmartConnectNoSSL, Disconnect
from pyVmomi import vim
import atexit
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def WaitForTask(task):
    """
    同步等待vCenter任务完成
    :param task: vCenter任务对象
    :raises: 任务失败时抛出异常
    """
    while task.info.state not in [vim.TaskInfo.State.success, vim.TaskInfo.State.error]:
        pass
    if task.info.state == vim.TaskInfo.State.error:
        raise task.info.error

async def migrate_vm_advanced(vm_names_or_source_host_ip, target_host_ip: str) -> dict:
    """
    高级迁移工具，支持：
    - 传入单个虚拟机名称字符串，迁移该虚拟机到目标主机
    - 传入多个虚拟机名称列表，迁移这些虚拟机到目标主机
    - 传入宿主机IP字符串，批量迁移该宿主机下所有虚拟机到目标主机

    :param vm_names_or_source_host_ip: 单个虚拟机名称（str）、多个虚拟机名称列表（list[str]）或宿主机IP（str）
    :param target_host_ip: 目标宿主机IP地址
    :return: 迁移结果状态和信息
    """
    try:
        # 如果传入的是字符串，且包含换行符或逗号，拆分为列表
        if isinstance(vm_names_or_source_host_ip, str) and ("\n" in vm_names_or_source_host_ip or "," in vm_names_or_source_host_ip):
            # 支持换行或逗号分隔
            if "\n" in vm_names_or_source_host_ip:
                vm_names_or_source_host_ip = [name.strip() for name in vm_names_or_source_host_ip.split("\n") if name.strip()]
            else:
                vm_names_or_source_host_ip = [name.strip() for name in vm_names_or_source_host_ip.split(",") if name.strip()]

        vcenter_host = os.getenv("VCENTER_HOST")
        vcenter_user = os.getenv("VCENTER_USER")
        vcenter_password = os.getenv("VCENTER_PASSWORD")

        si = SmartConnectNoSSL(host=vcenter_host, user=vcenter_user, pwd=vcenter_password, port=443)
        atexit.register(Disconnect, si)

        content = si.RetrieveContent()

        target_host = content.searchIndex.FindByIp(None, target_host_ip, False)
        if not target_host:
            return {"status": "error", "message": f"未能找到目标宿主机IP地址为 {target_host_ip}"}

        # 判断传入参数类型
        if isinstance(vm_names_or_source_host_ip, list):
            # 多个虚拟机名称列表，迁移指定虚拟机
            vm_names = vm_names_or_source_host_ip
            container = content.viewManager.CreateContainerView(content.rootFolder, [vim.VirtualMachine], True)
            vms = container.view
            vm_map = {vm.name: vm for vm in vms}
            selected_vms = []
            for name in vm_names:
                if name in vm_map:
                    selected_vms.append(vm_map[name])
                else:
                    return {"status": "error", "message": f"未能找到名称为 {name} 的虚拟机"}

            async def migrate_single_vm(vm):
                try:
                    reloc_spec = vim.VirtualMachineRelocateSpec()
                    reloc_spec.host = target_host
                    task = vm.RelocateVM_Task(reloc_spec)
                    await asyncio.get_event_loop().run_in_executor(None, WaitForTask, task)
                    return {"vm_name": vm.name, "status": "success", "message": f"虚拟机 {vm.name} 迁移成功"}
                except Exception as e:
                    return {"vm_name": vm.name, "status": "error", "message": str(e)}

            tasks = [migrate_single_vm(vm) for vm in selected_vms]
            results = await asyncio.gather(*tasks)

            return {"status": "success", "message": f"迁移完成，共迁移虚拟机数量: {len(selected_vms)}", "results": results}

        elif isinstance(vm_names_or_source_host_ip, str):
            # 先判断是否为宿主机IP
            source_host = content.searchIndex.FindByIp(None, vm_names_or_source_host_ip, False)
            if source_host:
                # 批量迁移宿主机下所有虚拟机
                vms = source_host.vm
                if not vms:
                    return {"status": "error", "message": f"宿主机 {vm_names_or_source_host_ip} 上没有虚拟机"}

                async def migrate_single_vm(vm):
                    try:
                        reloc_spec = vim.VirtualMachineRelocateSpec()
                        reloc_spec.host = target_host
                        task = vm.RelocateVM_Task(reloc_spec)
                        await asyncio.get_event_loop().run_in_executor(None, WaitForTask, task)
                        return {"vm_name": vm.name, "status": "success", "message": f"虚拟机 {vm.name} 迁移成功"}
                    except Exception as e:
                        return {"vm_name": vm.name, "status": "error", "message": str(e)}

                tasks = [migrate_single_vm(vm) for vm in vms]
                results = await asyncio.gather(*tasks)

                return {"status": "success", "message": f"批量迁移完成，共迁移虚拟机数量: {len(vms)}", "results": results}
            else:
                # 迁移单台虚拟机
                container = content.viewManager.CreateContainerView(content.rootFolder, [vim.VirtualMachine], True)
                vms = container.view
                vm = None
                for vm_instance in vms:
                    if vm_instance.name == vm_names_or_source_host_ip:
                        vm = vm_instance
                        break

                if not vm:
                    return {"status": "error", "message": f"未能找到名称为 {vm_names_or_source_host_ip} 的虚拟机"}

                reloc_spec = vim.VirtualMachineRelocateSpec()
                reloc_spec.host = target_host

                task = vm.RelocateVM_Task(reloc_spec)

                await asyncio.get_event_loop().run_in_executor(None, WaitForTask, task)

                return {"status": "success", "message": f"虚拟机 {vm.name} 已成功迁移到主机 {target_host.name}"}
        else:
            return {"status": "error", "message": "参数类型错误，必须是字符串或字符串列表"}

    except Exception as e:
        return {"status": "error", "message": str(e)}

